from nl2sql.model.llm import LLM
from nl2sql.engines.text2sql.prompt import prompt


class SQLGeneratorAgent:
    def __init__(self,
                 llm: LLM) -> None:
        self.llm = llm

    def text2sql(self,
                 query: str,
                 data_info: str,
                 evidence: str) -> str:

        filled_prompt = prompt.format(user_query=query,
                                      data_info=data_info,
                                      evidence=evidence)

        resp = self.llm.chat(prompt=filled_prompt)
        return resp

